import argparse
import ast

Dnum = 0
Rnum = 0
# 示例代码
code1 = """
#初始化最高温度
maxTemp=33
#初始化温度计
thermometer=Device()
#获取温度
temperature=thermometer.Get_Temperature()
#初始化空调
air_cond=Device()
#设置空调开关
#air_cond.On=0
a=a+b
if temperature>25:
    air_cond.On=1
    maxTemp=22
else:
    air_cond.On=2
    maxTemp=44
    
for i in range(3):
    print(i)
while a < 10:
    a += 1
"""


class registers:
    rx = None
    name = None

    def __init__(self, rx, name):
        self.rx = rx
        self.name = name


AssCode = []  # 汇编代码
# 用于存储已使用的寄存器
used_registers = {
    "d0": 0,
    "d1": 0,
    "d2": 0,
    "d3": 0,
    "d4": 0,
    "d5": 0,

    "r0": 0,
    "r1": 0,
    "r2": 0,
    "r3": 0,
    "r4": 0,
}
name_register = []


def get_available_register(type, name=None):
    '''
    分配寄存器
    :param type: 请求的寄存器的类型
    :return:
    '''
    if type == "d":
        for i in range(0, 5):  # 从 d1 开始尝试分配
            reg = f"d{i}"
            if not used_registers[reg]:
                used_registers[reg] = 1
                name_register.append(registers(reg, name))
                return reg
    elif type == "r":
        for i in range(1, 16):  # 从 r1 开始尝试分配
            reg = f"r{i}"
            if not used_registers[reg]:
                used_registers[reg] = 1
                name_register.append(registers(reg, name))
                return reg
    raise Exception("寄存器资源已耗尽")


import ast
import sys


class InstructionGenerator(ast.NodeVisitor):
    def __init__(self):
        self.instructions = []
        self.variable_map = {}
        self.register_id = 0
        self.degister_id = 0
        self.label_id = 0
        self.variable_use={}

    def get_available_register(self, type, num=-1):
        if type == "d" and num == -1:  # num=-1，说明未指定寄存器的值
            num = self.degister_id
            register = f"{type}{num}"
            self.degister_id += 1
        elif type == "r" and num == -1:
            num = self.register_id
            register = f"{type}{num}"
            self.register_id += 1
        elif type == "d" and num != -1:
            register = f"{type}{num}"
        elif type == "r" and num != -1:
            register = f"{type}{num}"
        return register

    def visit_BinOp(self, node):
        if isinstance(node.left, ast.Name):
            left = self.variable_map.get(node.left.id, node.left.id)
        else:
            left = node.left.n if isinstance(node.left, ast.Num) else self.visit(node.left)

        if isinstance(node.right, ast.Name):
            right = self.variable_map.get(node.right.id, node.right.id)
        else:
            right = node.right.n if isinstance(node.right, ast.Num) else self.visit(node.right)

        target_register = self.get_available_register('r')

        if isinstance(node.op, ast.Add):
            self.instructions.append(f"add {target_register} {left} {right}")
        elif isinstance(node.op, ast.Sub):
            self.instructions.append(f"sub {target_register} {left} {right}")
        elif isinstance(node.op, ast.Mult):
            self.instructions.append(f"mul {target_register} {left} {right}")
        elif isinstance(node.op, ast.Div):
            self.instructions.append(f"div {target_register} {left} {right}")
        elif isinstance(node.op, ast.Mod):
            self.instructions.append(f"mod {target_register} {left} {right}")
        else:
            raise ValueError(f"Unsupported binary operation: {type(node.op)}")

        self.generic_visit(node)
        return target_register

    def visit_Assign(self, node):
        if isinstance(node.targets[0], ast.Name):
            target = node.targets[0].id
            if isinstance(node.value, ast.Num):
                value = node.value.n
                self.variable_use[target] = value
                self.instructions.append(f"define {target} {value}")
            elif isinstance(node.value, ast.Name):
                value = node.value.id
                self.instructions.append(f"alias {target} {value}")
            elif isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Attribute):
                if node.value.func.attr.startswith("Get_"):
                    if isinstance(node.value.func.value, ast.Name):
                        device_name = node.value.func.value.id
                        attr_name = node.value.func.attr[4:].capitalize()  # 提取 'Get_' 后面的部分
                        self.instructions.append(f"l {target} {device_name} {attr_name}")
            elif isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
                if node.value.func.id == "Device":
                    register = self.get_available_register(type='d')
                    self.instructions.append(f"alias {target} {register}")
            else:
                value = self.visit(node.value)
                target_register = self.get_available_register('r')
                self.instructions.append(f"alias {target} {target_register}")
                self.variable_map[target] = target_register
        elif isinstance(node.targets[0], ast.Attribute):
            target = node.targets[0].value.id + "." + node.targets[0].attr
            value = node.value.n
            self.instructions.append(f"s {target.split('.')[0]} {target.split('.')[1]} {value}")
        self.generic_visit(node)

    def visit_Expr(self, node):
        if isinstance(node.value, ast.Call) and isinstance(node.value.func, ast.Name):
            if node.value.func.id == "Device":
                register = self.get_available_register(type='d', name=node.value.func.id)
                self.instructions.append(f"alias {node.value.func.id} {register}")
        self.generic_visit(node)

    def visit_Compare(self, node,  target_label=None):
        if isinstance(node.left, ast.Name):
            left = self.variable_map.get(node.left.id, node.left.id)
        else:
            left = node.left.n if isinstance(node.left, ast.Num) else self.visit(node.left)

        if isinstance(node.comparators[0], ast.Name):
            right = self.variable_map.get(node.comparators[0].id, node.comparators[0].id)
        else:
            right = node.comparators[0].n
        if target_label is None:
            target_label = f"label{self.label_id}:"
            self.label_id += 1

        if isinstance(node.ops[0], ast.Eq):
            self.instructions.append(f"bne {left} {right} {target_label[:-1]}")
        elif isinstance(node.ops[0], ast.NotEq):
            self.instructions.append(f"bne {left} {right} {target_label[:-1]}")
        elif isinstance(node.ops[0], ast.Lt):
            self.instructions.append(f"blt {left} {right} {target_label[:-1]}")
        elif isinstance(node.ops[0], ast.LtE):
            self.instructions.append(f"ble {left} {right} {target_label[:-1]}")
        elif isinstance(node.ops[0], ast.Gt):
            self.instructions.append(f"bgt {left} {right} {target_label[:-1]}")
        elif isinstance(node.ops[0], ast.GtE):
            self.instructions.append(f"bge {left} {right} {target_label[:-1]}")
        #self.instructions.append(f" {target_label} ")
        self.generic_visit(node)

    def visit_Call(self, node):
        if isinstance(node.func, ast.Attribute):
            if node.func.attr.startswith("Get_"):
                pass
                # if isinstance(node.func.value, ast.Name):
                #     device_name = node.func.value.id
                #     attr_name = node.func.attr[4:].capitalize()  # 提取 'Get_' 后面的部分
                #     target_register = self.get_available_register('r')
                #     self.instructions.append(f"l {target_register} {device_name} {attr_name}")
                #     return target_register
                # else:
                #     self.visit(node.func.value)
                #     attr_name = node.func.attr[4:].capitalize()  # 提取 'Get_' 后面的部分
                #     target_register = self.get_available_register('r')
                #     self.instructions.append(f"l {target_register} r{self.register_id - 1} {attr_name}")
                #     return target_register
        elif isinstance(node.func, ast.Name):
            if node.func.id == "log":
                if isinstance(node.args[0], ast.Name):
                    arg = self.variable_map.get(node.args[0].id, node.args[0].id)
                else:
                    arg = node.args[0].n
                target_register = self.get_available_register('r')
                self.instructions.append(f"log {target_register} {arg}")
                return target_register
            elif node.func.id == "sqrt":
                if isinstance(node.args[0], ast.Name):
                    arg = self.variable_map.get(node.args[0].id, node.args[0].id)
                else:
                    arg = node.args[0].n
                target_register = self.get_available_register('r')
                self.instructions.append(f"sqrt {target_register} {arg}")
                return target_register
            elif node.func.id == "exp":
                if isinstance(node.args[0], ast.Name):
                    arg = self.variable_map.get(node.args[0].id, node.args[0].id)
                else:
                    arg = node.args[0].n
                target_register = self.get_available_register('r')
                self.instructions.append(f"exp {target_register} {arg}")
                return target_register
            elif node.func.id == "abs":
                if isinstance(node.args[0], ast.Name):
                    arg = self.variable_map.get(node.args[0].id, node.args[0].id)
                else:
                    arg = node.args[0].n
                target_register = self.get_available_register('r')
                self.instructions.append(f"abs {target_register} {arg}")
                return target_register
        self.generic_visit(node)

    def visit_If(self, node):
        end_label = f"endLabelIf{self.label_id}:"
        self.label_id += 1

        if isinstance(node.test, ast.Compare):
            if len(node.test.ops) > 0:
                op = node.test.ops[0]
                left = node.test.left
                right = node.test.comparators[0]

                # 获取左值
                if isinstance(left, ast.Name):
                    left_val = self.variable_map.get(left.id, left.id)
                else:
                    left_val = left.n if isinstance(left, ast.Num) else self.visit(left)

                # 获取右值
                if isinstance(right, ast.Name):
                    right_val = self.variable_map.get(right.id, right.id)
                else:
                    right_val = right.n if isinstance(right, ast.Num) else self.visit(right)

                # 根据不同的比较操作符生成对应的指令
                if isinstance(op, ast.Eq):#==
                    self.instructions.append(f"bne {left_val} {right_val} {end_label[:-1]}")
                elif isinstance(op, ast.NotEq):#!=
                    self.instructions.append(f"beq {left_val} {right_val} {end_label[:-1]}")
                elif isinstance(op, ast.Lt):#<
                    self.instructions.append(f"bgt {left_val} {right_val} {end_label[:-1]}")
                elif isinstance(op, ast.LtE):#<=
                    self.instructions.append(f"bge {left_val} {right_val} {end_label[:-1]}")
                elif isinstance(op, ast.Gt):#>
                    self.instructions.append(f"blt {left_val} {right_val} {end_label[:-1]}")
                elif isinstance(op, ast.GtE):#>=
                    self.instructions.append(f"ble {left_val} {right_val} {end_label[:-1]}")
        else:
            self.visit(node.test)
            self.instructions.append(f"xxxxxxxxxx")

        if len(node.body) > 0:
            for stmt in node.body:
                self.visit(stmt)
        self.instructions.append(f"breqz 0 {len(node.orelse)+2}")
        self.instructions.append(f" {end_label} ")
        for stmt in node.orelse:
            self.visit(stmt)


    def visit_While(self, node):
        """
        while a==b:
           主体
        ->
        label_0:
        beq a b label_1#如果a=b，跳转到label_1
        主体
        j label_0
        label_1:
        :param node:
        :return:
        """
        loop_start = f"label{self.label_id}:"
        self.label_id += 1
        loop_end = f"label{self.label_id}:"
        self.label_id += 1
        self.instructions.append(f" {loop_start} ")#  开始标签
        self.visit_Compare(node.test, loop_end)#判断跳出
        for stmt in node.body:#主体
            self.visit(stmt)
        self.instructions.append(f"j {loop_start[:-1]} ")#跳转到开始标签

        self.instructions.append(f" {loop_end} ")

    def visit_FunctionDef(self, node):
        self.instructions.append(f" func_{node.name} ")
        for arg in node.args.args:
            self.visit(arg)
        for stmt in node.body:
            self.visit(stmt)
        self.instructions.append(f" end_func_{node.name} ")

    def generate(self, code):
        tree = ast.parse(code)
        self.visit(tree)
        return self.instructions

print("使用方法：python transformer.py --path my_code.py")
print("或者直接在IDE运行")
parser = argparse.ArgumentParser()
parser.add_argument('--path', type=str, default='testPy.py')  # 默认文件路径为 testPy.py
args = parser.parse_args()

with open(args.path, "r", encoding="utf-8") as f:
    code = f.read()
    tree = ast.parse(code)

print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
print("汇编码")
generator = InstructionGenerator()
instructions = generator.generate(code)

for instruction in instructions:
    print(str(instruction).strip())
# with open("testPy.py", "r", encoding="utf-8") as f:
#     code = f.read()
#     tree = ast.parse(code)
#
# print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
# print("汇编码")
# generator = InstructionGenerator()
# instructions = generator.generate(code)
#
# for instruction in instructions:
#     print(str(instruction).strip())
